import datetime
import os

import cv2
import numpy as np
from torch import nn, optim
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from data import YellowDataset
from net import Net
from PIL import Image, ImageDraw
from torch.nn.functional import one_hot

DEVICE = 'cpu'


class Train:
    def __init__(self, root, weight_path):
        self.summaryWriter = SummaryWriter('logs')
        self.train_dataset = YellowDataset(root,True)
        self.test_dataset = YellowDataset(root,False)
        self.train_dataloader = DataLoader(self.train_dataset, batch_size=50, shuffle=True)
        self.test_dataloader = DataLoader(self.test_dataset, batch_size=1, shuffle=True)
        self.net = Net().to(DEVICE)
        if os.path.exists(weight_path):
            self.net.load_state_dict(torch.load(weight_path))

        self.opt = optim.Adam(self.net.parameters())
        self.label_loss = nn.BCEWithLogitsLoss()

        self.position_loss = nn.MSELoss()
        self.sort_loss = nn.CrossEntropyLoss()

        self.train = False
        self.test = True

    def __call__(self):
        index1, index2 = 0, 0
        for epoch in range(1000):
            if self.train:

                for i, (img, label, position, sort) in enumerate(self.train_dataloader):
                    img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                    out_label, out_position, out_sort = self.net(img)
                    position_loss = self.position_loss(out_position, position)
                    out_sort=out_sort[torch.where(sort>=0)]
                    sort=sort[torch.where(sort>=0)]
                    sort_loss = self.sort_loss(out_sort, sort)
                    label_loss = self.label_loss(out_label, label)

                    train_loss = 0.2 * label_loss + position_loss * 0.6 + 0.2 * sort_loss
                    # print('train_sort:', sort)
                    # print('train_out_sort:', torch.argmax(torch.softmax(out_sort, dim=1), dim=1))
                    self.opt.zero_grad()
                    train_loss.backward()
                    self.opt.step()
                    if i % 10 == 0:
                        print(f'train_loss {i}===>>', train_loss.item())
                        self.summaryWriter.add_scalar('train_loss', train_loss, index1)
                        index1 += 1
                date_time = str(datetime.datetime.now()).replace(' ', '-').replace(':', '_').replace('.', '_')
                torch.save(self.net.state_dict(), f'param/{date_time}-{epoch}.pt')
                sum_acc = 0

                for i, (img, label, position, sort) in enumerate(self.test_dataloader):
                    img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                    out_label, out_position, out_sort = self.net(img)
                    out_sort = out_sort[torch.where(sort >= 0)]
                    sort = sort[torch.where(sort >= 0)]
                    position_loss = self.position_loss(out_position, position)
                    sort_loss = self.sort_loss(out_sort, sort)
                    label_loss = self.label_loss(out_label, label)
                    test_loss = 0.2 * label_loss + position_loss * 0.6 + 0.2 * sort_loss
                    # print('test_sort:', sort)
                    # print('test_out_sort:', torch.argmax(torch.softmax(out_sort, dim=1), dim=1))
                    sort_acc = torch.mean(torch.eq(sort, torch.argmax(torch.softmax(out_sort, dim=1), dim=1)).float())
                    sum_acc += sort_acc
                    if i % 5 == 0:
                        print(f'test_loss {i}===>>', test_loss.item())
                        self.summaryWriter.add_scalar('test_loss', test_loss, index2)
                        index2 += 1
                avg_acc = sum_acc / i
                print(f'sort_acc {i}==>>', avg_acc)
                self.summaryWriter.add_scalar('avg_acc', avg_acc, epoch)

            if self.test:
                sum_acc = 0
                i =0
                filename = './data/newTest/'
                for i, (img, label, position, sort) in enumerate(self.test_dataloader):
                    img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                    out_label, out_position, out_sort = self.net(img)

                    out_label = torch.sigmoid(out_label)

                    position = position * 300
                    position = [int(i) for i in position[0]]
                    out_position = out_position * 300
                    out_position = [int(i) for i in out_position[0]]
                    out_sort = torch.argmax(torch.softmax(out_sort, dim=1))
                    new_img = torch.squeeze(img)
                    new_img = new_img.permute(1, 2, 0)
                    new_img = np.array(new_img.cpu())
                    new_img = new_img.copy()

                    cv2.rectangle(new_img, (position[0], position[1]), (position[2], position[3]), (0, 255, 0), 3)
                    cv2.putText(new_img, str(sort.item()), (position[0], position[1] - 3),
                                cv2.FONT_HERSHEY_SIMPLEX,
                                1, (0, 255, 0), 3)
                    if out_label.item() > 0.5:
                        cv2.rectangle(new_img, (out_position[0], out_position[1]), (out_position[2], out_position[3]),
                                      (0, 0, 255), 3)
                        cv2.putText(new_img, str(out_sort.item()), (out_position[0], out_position[1] - 3),
                                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 3)
                    cv2.imshow('new_img', new_img)
                    # nerdir = filename+str(i)+'.jpg'
                    # cv2.imwrite(nerdir,new_img)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()
                    i = i +1


if __name__ == '__main__':
    dataSet = 'D:\BaiduNetdiskWorkspace\homework\cv\pytorch-Single-target-detection-Minions\yellow_demo\yellow_demo\data'
    train = Train(dataSet,'')
    train()
